export class MCTSNode {
  constructor(state, parent = null) {
    this.state = state;
    this.parent = parent;
    this.children = [];
    this._results = new Map();
    this._N = 0; // 访问次数
    this.untried_actions = this.state.get_legal_actions();
  }

  get Q() {
    // 收益
    // 平局未处理
    let res = 0;
    let wins = this._results.get(this.parent.state.to_move) || 0; // 父节点的toMove就是本节点
    let loses = this._results.get(-this.parent.state.to_move) || 0;
    res = wins - loses;
    return res;
  }

  get N() {
    return this._N;
  }

  is_terminal_node() {
    return this.state.is_gameover();
  }

  is_full_expanded() {
    return this.untried_actions.length === 0;
  }

  expand() {
    const action = this.untried_actions.pop();

    const tState = this.state.chessmove(action);

    const tNode = new MCTSNode(tState, this); // 当前局面先走一步

    this.children.push(tNode);

    return tNode;
  }

  rollout() {
    // 走一步之后随机走子模拟到结束
    let curState = this.state;
    while (!curState.is_gameover()) {
      const possibleMoves = curState.get_legal_actions();
      const action =
        possibleMoves[Math.floor(Math.random() * possibleMoves.length)];
      curState = curState.chessmove(action);
    }

    return curState.game_result;
  }

  back(reward) {
    this._N++;
    if (this._results.has(reward)) {
      this._results.set(reward, this._results.get(reward) + 1); // 简单的把每种结局统计次数
    } else {
      this._results.set(reward, 1);
    }

    if (this.parent) {
      this.parent.back(reward);
    }
  }

  best_child(c_param = 1 / Math.sqrt(2)) {
    const choicesWeight = this.children.map((child) => {
      return (
        child.Q / child.N +
        c_param * Math.sqrt((2 * Math.log(this.N)) / child.N)
      );
    }); // UCB
    const maxChoiceIndex = choicesWeight.indexOf(Math.max(...choicesWeight));
    return this.children[maxChoiceIndex];
  }
}
